#include "ompl_extension/OctomapCostEstimator.h"

namespace ompl
{
  namespace base
  {

    void worldToGrid(const btVector3& origin, double wx, double wy, double wz, double resolution, int& gx, int& gy, int& gz)
    {
      gx = (int)((wx - origin.x()) * (1.0 / resolution));
      gy = (int)((wy - origin.y()) * (1.0 / resolution));
      gz = (int)((wz - origin.z()) * (1.0 / resolution));
    }

    void gridToWorld(const btVector3& origin, int gx, int gy, int gz, double resolution, double& wx, double& wy, double& wz)
    {
      wx = gx * resolution + origin.x();
      wy = gy * resolution + origin.y();
      wz = gz * resolution + origin.z();
    }

    void getVoxelsInBody(const bodies::Body& body, double resolution, std::vector<btVector3>& voxels)
    {
      bodies::BoundingSphere bounding_sphere;

      body.computeBoundingSphere(bounding_sphere);

      int x, y, z, x_min, x_max, y_min, y_max, z_min, z_max;
      double xw, yw, zw;
      btVector3 v;

      worldToGrid(bounding_sphere.center, bounding_sphere.center.x() - bounding_sphere.radius,
                                          bounding_sphere.center.y() - bounding_sphere.radius,
                                          bounding_sphere.center.z() - bounding_sphere.radius,
                                          resolution,
                                          x_min, y_min, z_min);
      worldToGrid(bounding_sphere.center, bounding_sphere.center.x() + bounding_sphere.radius,
                                          bounding_sphere.center.y() + bounding_sphere.radius,
                                          bounding_sphere.center.z() + bounding_sphere.radius,
                                          resolution,
                                          x_max, y_max, z_max);

      for(x = x_min; x <= x_max; ++x)
      {
        for(y = y_min; y <= y_max; ++y)
        {
          for(z = z_min; z <= z_max; ++z)
          {
            gridToWorld(bounding_sphere.center, x, y, z, resolution, xw, yw, zw);
            if(body.containsPoint(xw, yw, zw))
            {
              v.setX(xw);
              v.setY(yw);
              v.setZ(zw);
              voxels.push_back(v);
            }
          }
        }
      }
    }


    OctomapStateCostEstimator::OctomapStateCostEstimator(ompl::base::SpaceInformation* si,
                                                         planning_environment::CollisionModelsInterface* cmi,
                                                         const std::string& group_name,
                                                         const ros::NodeHandle& node_handle,
                                                         double resolution) : ompl::base::StateCostEstimator(si),
                                                         octree_(resolution)
    {
      resolution_ = resolution;

      node_handle_ = node_handle;

      ros::service::waitForService("/octomap_binary");
      const static std::string servname = "octomap_binary";
      octomap_ros::GetOctomap::Request req;
      octomap_ros::GetOctomap::Response resp;

      if(node_handle_.ok() && ros::service::call(servname, req, resp))
      {
        octomap_ros::OctomapBinary map_msg = resp.map;
        octomap::octomapMsgToMap(map_msg, octree_);
      }
      else
      {
        ROS_ERROR("Failed to call service /octomap_binary");
      }

      cmi_ = cmi;

      joint_group_ = cmi_->getKinematicModel()->getModelGroup(group_name);

      group_links_ = joint_group_->getGroupLinkModels();

      if(!ompl_ros_interface_helper::jointGroupToOmplStateSpacePtr(joint_group_, ompl_state_to_kinematic_state_mapping_, kinematic_state_to_ompl_state_mapping_))
      {
        ROS_ERROR("Could not set up the ompl state space");
        return;
      }

      // sample points in each link

      for(unsigned int i = 0; i < group_links_.size(); ++i)
      {
        const planning_models::KinematicModel::LinkModel* link = group_links_[i];
        if(link)
        {
          bodies::Body* body = bodies::createBodyFromShape(link->getLinkShape());

          std::vector<btVector3> voxels;
          getVoxelsInBody(*body, resolution_, voxels);

          link_samples_[link] = voxels;
        }
      }
    }


    double OctomapStateCostEstimator::computeCost(const ompl::base::State* state) const
    {
      planning_models::KinematicState* kinematic_state = cmi_->getPlanningSceneState();
      if(!kinematic_state)
      {
        ROS_ERROR_STREAM("Planning scene hasn't been set");
        return 0;
      }

      // set state
      planning_models::KinematicState::JointStateGroup* joint_state_group = kinematic_state->getJointStateGroup(joint_group_->getName());
      omplStateToKinematicStateGroup(state, ompl_state_to_kinematic_state_mapping_, joint_state_group);
      joint_state_group->updateKinematicLinks();

      std::map<const planning_models::KinematicModel::LinkModel*, std::vector<btVector3> > transformed_link_samples;

      for(unsigned int i = 0; i < group_links_.size(); ++i)
      {
        const planning_models::KinematicModel::LinkModel* link = group_links_[i];
        const planning_models::KinematicState::LinkState* link_state = kinematic_state->getLinkState(link->getName());

        if(link)
        {

          btTransform trans = link_state->getGlobalCollisionBodyTransform();
          std::map<const planning_models::KinematicModel::LinkModel*, std::vector<btVector3> >::const_iterator find_it = link_samples_.find(link);

          transformed_link_samples[link] = find_it->second;
          for(unsigned int j = 0; j < find_it->second.size(); ++j)
          {
            transformed_link_samples[link][j] = trans * find_it->second[j];
          }
        }
      }

      double cost = 0;

      for(std::map<const planning_models::KinematicModel::LinkModel*, std::vector<btVector3> >::const_iterator it = transformed_link_samples.begin();
          it != transformed_link_samples.end();
          ++it)
      {
        for(unsigned int i = 0; i < it->second.size(); ++i)
        {
          octomap::OcTreeNode* node = octree_.search(it->second[i].x(), it->second[i].y(), it->second[i].z());
          if(node)
          {
            cost += node->getOccupancy();
          }
        }
      }

      return cost;
    }




    OctomapMotionCostEstimator::OctomapMotionCostEstimator(ompl::base::SpaceInformation* si,
                                                           planning_environment::CollisionModelsInterface* cmi,
                                                           const std::string& group_name,
                                                           const ros::NodeHandle& node_handle,
                                                           double resolution,
                                                           unsigned int interpolate_num) : ompl::base::MotionCostEstimator(si),
                                                           octree_(resolution), interpolate_num_(interpolate_num)
    {
      resolution_ = resolution;

      node_handle_ = node_handle;

      ros::service::waitForService("/octomap_binary");
      const static std::string servname = "octomap_binary";
      octomap_ros::GetOctomap::Request req;
      octomap_ros::GetOctomap::Response resp;

      if(node_handle_.ok() && ros::service::call(servname, req, resp))
      {
        octomap_ros::OctomapBinary map_msg = resp.map;
        octomap::octomapMsgToMap(map_msg, octree_);
      }
      else
      {
        ROS_ERROR("Failed to call service /octomap_binary");
      }

      cmi_ = cmi;

      joint_group_ = cmi_->getKinematicModel()->getModelGroup(group_name);

      group_links_ = joint_group_->getGroupLinkModels();

      if(!ompl_ros_interface_helper::jointGroupToOmplStateSpacePtr(joint_group_, ompl_state_to_kinematic_state_mapping_, kinematic_state_to_ompl_state_mapping_))
      {
        ROS_ERROR("Could not set up the ompl state space");
        return;
      }

      // sample points in each link

      for(unsigned int i = 0; i < group_links_.size(); ++i)
      {
        const planning_models::KinematicModel::LinkModel* link = group_links_[i];
        if(link)
        {
          bodies::Body* body = bodies::createBodyFromShape(link->getLinkShape());

          std::vector<btVector3> voxels;
          getVoxelsInBody(*body, resolution_, voxels);

          link_samples_[link] = voxels;
        }
      }
    }


    double OctomapMotionCostEstimator::computeCost(const ompl::base::State* s1, const ompl::base::State* s2) const
    {
      planning_models::KinematicState* kinematic_state = cmi_->getPlanningSceneState();
      if(!kinematic_state)
      {
        ROS_ERROR_STREAM("Planning scene hasn't been set");
        return 0;
      }


      double sub_length = si_->distance(s1, s2) / interpolate_num_;
      std::vector<double> endpoint_costs(interpolate_num_ + 1);

      planning_models::KinematicState::JointStateGroup* joint_state_group = kinematic_state->getJointStateGroup(joint_group_->getName());
      base::State* interp_state = si_->allocState();

      for(unsigned int n = 0; n <= interpolate_num_; ++n)
      {
        if(n == 0)
          si_->copyState(interp_state, s1);
        else if(n == interpolate_num_)
          si_->copyState(interp_state, s2);
        else
        {
          si_->getStateSpace()->interpolate(s1, s2, n / (double)interpolate_num_, interp_state);
        }

        omplStateToKinematicStateGroup(interp_state, ompl_state_to_kinematic_state_mapping_, joint_state_group);
        joint_state_group->updateKinematicLinks();

        std::map<const planning_models::KinematicModel::LinkModel*, std::vector<btVector3> > transformed_link_samples;

        for(unsigned int i = 0; i < group_links_.size(); ++i)
        {
          const planning_models::KinematicModel::LinkModel* link = group_links_[i];
          const planning_models::KinematicState::LinkState* link_state = kinematic_state->getLinkState(link->getName());

          if(link)
          {

            btTransform trans = link_state->getGlobalCollisionBodyTransform();
            std::map<const planning_models::KinematicModel::LinkModel*, std::vector<btVector3> >::const_iterator find_it = link_samples_.find(link);

            transformed_link_samples[link] = find_it->second;
            for(unsigned int j = 0; j < find_it->second.size(); ++j)
            {
              transformed_link_samples[link][j] = trans * find_it->second[j];
            }
          }
        }


        double cost_ = 0.0;

        for(std::map<const planning_models::KinematicModel::LinkModel*, std::vector<btVector3> >::const_iterator it = transformed_link_samples.begin();
            it != transformed_link_samples.end();
            ++it)
        {
          for(unsigned int i = 0; i < it->second.size(); ++i)
          {
            octomap::OcTreeNode* node = octree_.search(it->second[i].x(), it->second[i].y(), it->second[i].z());
            if(node)
            {
              cost_ += node->getOccupancy();
            }
          }
        }

        endpoint_costs[n] = cost_;
      }


      double cost = 0.0;

      for(unsigned int i = 0; i < interpolate_num_; ++i)
      {
        cost += sub_length * (1 + 0.5 * (endpoint_costs[i] + endpoint_costs[i + 1]));
      }


      return cost;
    }


  }
}
